#include "fdmU2d.h"

#include "../global.h"
#include <cstdio>
#include <iostream>
#include <cstring>
#include "diff2d.h"
#include "initVars.h"
#include "delVar.h"

using namespace std;

int fdmU2d::iter(diff2D::diffList diffFun,int xNum,int yNum,double **x0,double **b,double **x1)
{
    switch(schemeNum){
        case 1:
            return gaussSeidel(diffFun,xNum,yNum,x0,b,x1);
        case 2:
            return jaccobi(diffFun,xNum,yNum,x0,b,x1);
        case 3:
            return sor(diffFun,xNum,yNum,x0,b,x1);
        case 4:
            if(fdmU2d::conjGradIter.inited==0){
                fdmU2d::conjGradIter.init(diffFun,xNum,yNum,x0,b,x1);
            }
            return fdmU2d::conjGradIter.iter(diffFun,xNum,yNum,x0,b,x1);
    }
    
    return -1;
}

int fdmU2d::jaccobi(diff2D::diffList diffFun,int xNum,int yNum,double **x0,double **b,double **x1)
{
    double residual=0;
    double temp;
    for(int k=1;k<xNum-1;k++){
        for(int p=1;p<yNum-1;p++){
            x1[k][p]=b[k][p];
            for(int q=0;q<diffFun.lNum;q++){
                x1[k][p]=x1[k][p]-x0[k+diffFun.lList[q].i][p+diffFun.lList[q].j]*diffFun.lCoef[q];
            }
            for(int q=0;q<diffFun.uNum;q++){
                x1[k][p]=x1[k][p]-x0[k+diffFun.uList[q].i][p+diffFun.uList[q].j]*diffFun.uCoef[q];
            }
            temp=abs(x1[k][p]-x0[k][p]);
            if(temp>residual){
                residual=temp;
            }
        }
    }
    cout<<"\tResidual="<<residual<<endl;
    if(residual>resMin)
        return 0;
    return 1;
}

int fdmU2d::gaussSeidel(diff2D::diffList diffFun,int xNum,int yNum,double **x0,double **b,double **x1)
{
    double residual=0;
    double temp;
    // for(int p=0;p<yNum;p++){
    //     x1[0][p]=x0[0][p];
    //     x1[xNum-1][p]=x0[xNum-1][p];
    // }
    // for(int k=0;k<xNum;k++){
    //     x1[k][0]=x0[k][0];
    //     x1[k][yNum-1]=x0[k][yNum-1];
    // }
    for(int k=1;k<xNum-1;k++){
        for(int p=1;p<yNum-1;p++){
            x1[k][p]=b[k][p];
            for(int q=0;q<diffFun.lNum;q++){
                x1[k][p]=x1[k][p]-x1[k+diffFun.lList[q].i][p+diffFun.lList[q].j]*diffFun.lCoef[q];
            }
            for(int q=0;q<diffFun.uNum;q++){
                x1[k][p]=x1[k][p]-x0[k+diffFun.uList[q].i][p+diffFun.uList[q].j]*diffFun.uCoef[q];
            }
            temp=abs(x1[k][p]-x0[k][p]);
            if(temp>residual){
                residual=temp;
            }
        }
    }
    cout<<"\tResidual="<<residual<<endl;
    if(residual>resMin)
        return 0;
    return 1;
}

int fdmU2d::sor(diff2D::diffList diffFun,int xNum,int yNum,double **x0,double **b,double **x1)
{
    double residual=0;
    double temp;
    for(int k=1;k<xNum-1;k++){
        for(int p=1;p<yNum-1;p++){
            x1[k][p]=b[k][p];
            for(int q=0;q<diffFun.lNum;q++){
                x1[k][p]=x1[k][p]-x1[k+diffFun.lList[q].i][p+diffFun.lList[q].j]*diffFun.lCoef[q];
            }
            for(int q=0;q<diffFun.uNum;q++){
                x1[k][p]=x1[k][p]-x0[k+diffFun.uList[q].i][p+diffFun.uList[q].j]*diffFun.uCoef[q];
            }
            x1[k][p]=sorFac*(x1[k][p])+(1-sorFac)*(x0[k][p]);
            temp=abs(x1[k][p]-x0[k][p]);
            if(temp>residual){
                residual=temp;
            }
        }
    }
    cout<<"\tResidual="<<residual<<endl;
    if(residual>resMin)
        return 0;
    return 1;
}

void fdmU2d::conjGrad::init(diff2D::diffList diffFun,int xNum,int yNum,double **x0,double **b,double **x1)
{
    r0=zeros2D(xNum,yNum);
    r1=zeros2D(xNum,yNum);
    p0=zeros2D(xNum,yNum);
    p1=zeros2D(xNum,yNum);

    double **ax0=matVec(diffFun,x0,xNum,yNum);

    for(int k=0;k<xNum;k++){
        for(int p=0;p<yNum;p++){
            r0[k][p]=b[k][p]-ax0[k][p];
            p0[k][p]=r0[k][p];
        }
    }

    delMat2D(ax0,xNum,yNum);
    inited=1;
    return;
}

int fdmU2d::conjGrad::iter(diff2D::diffList diffFun,int xNum,int yNum,double **x0,double **b,double **x1)
{
    double residual=0;

    double alpha;
    double **ap0=matVec(diffFun,p0,xNum,yNum);
    alpha=dot(r0,r0,xNum,yNum)/dot(p0,ap0,xNum,yNum);
    
    for(int k=1;k<xNum-1;k++){
        for(int p=1;p<yNum-1;p++){
            x1[k][p]=x0[k][p]+alpha*p0[k][p];
            r1[k][p]=r0[k][p]-alpha*ap0[k][p];
            if(abs(r0[k][p])>residual){
                residual=abs(r0[k][p]);
            }
        }
    }

    double beta=0;
    beta=dot(r1,r1,xNum,yNum)/dot(r0,r0,xNum,yNum);
    for(int k=1;k<xNum;k++){
        for(int p=1;p<yNum;p++){
            p1[k][p]=r1[k][p]+beta*p0[k][p];
        }
    }

    updateVec(p0,p1,xNum,yNum);
    updateVec(r0,r1,xNum,yNum);
    delMat2D(ap0,xNum,yNum);
    cout<<"\tResidual="<<residual<<endl;
    if(residual>resMin)
        return 0;
    return 1;
}

int fdmU2d::copyBc(int xNum,int yNum,double **x0,double **x1)
{
    for(int p=0;p<yNum;p++){
        x1[0][p]=x0[0][p];
        x1[xNum-1][p]=x0[xNum-1][p];
    }
    for(int k=0;k<xNum;k++){
        x1[k][0]=x0[k][0];
        x1[k][yNum-1]=x0[k][yNum-1];
    }
    return 1;
}

int fdmU2d::readScheme(char *fileName)
{
    char filePath[256];
    strcpy(filePath,currentPath);
    strcat(filePath,fileName);
    cout<<"Reading scheme from "<<filePath<<endl;
    ifstream fin(filePath);
    streambuf *oldcin;
    oldcin = cin.rdbuf(fin.rdbuf());

    string tempstr;
    cin>>scheme;
    if(scheme=="gaussSeidel"){
        schemeNum=1;
        cout<<"\t Gauss-Seidel"<<endl;
    }
    else if(scheme=="jaccobi"){
        schemeNum=2;
        cout<<"\t Jaccobi"<<endl;
    }
    else if(scheme=="sor"){
        schemeNum=3;
        cout<<"\t SOR"<<endl;
        cin>>tempstr;
        if(tempstr=="sorFac"){
            cin>>sorFac;
        }
        cout<<"\t sorFac="<<sorFac<<endl;
    }
    else if(scheme=="conjugateGradient"){
        schemeNum=4;
        cout<<"\t Conjugate Gradient"<<endl;
    }
    
    cin>>tempstr;
    if(tempstr=="residual"){
        cin>>resMin;
    }

    cin.rdbuf(oldcin);
    fin.close();
    cout<<"Done reading scheme."<<endl;
    return 1;
}

double fdmU2d::dot(double **vec1,double **vec2,int xNum,int yNum)
{
    double sum=0;
    for(int k=0;k<xNum;k++){
        for(int p=0;p<yNum;p++){
            sum+=vec1[k][p]*vec2[k][p];
        }
    }
    return sum;
}

double **fdmU2d::matVec(diff2D::diffList diffFun,double **vec,int xNum,int yNum)
{
    double **ax0=zeros2D(xNum,yNum);

    for(int k=1;k<xNum-1;k++){
        for(int p=1;p<yNum-1;p++){
            ax0[k][p]=vec[k][p];
            for(int q=0;q<diffFun.lNum;q++){
                ax0[k][p]=ax0[k][p]+vec[k+diffFun.lList[q].i][p+diffFun.lList[q].j]*diffFun.lCoef[q];
            }
            for(int q=0;q<diffFun.uNum;q++){
                ax0[k][p]=ax0[k][p]+vec[k+diffFun.uList[q].i][p+diffFun.uList[q].j]*diffFun.uCoef[q];
            }
        }
    }

    return ax0;
}

int fdmU2d::updateVec(double **vec0,double **vec1,int xNum,int yNum)
{
    for(int k=0;k<xNum;k++){
        for(int p=0;p<yNum;p++){
            vec0[k][p]=vec1[k][p];
        }
    }
    return 1;
}

